#include "misc.h"

#include <algorithm>
#include <atomic>
#include <thread>
#include <vector>

namespace tinyndarray {

void ApplyOpSimple(NdArray& ret, UnaryElementOp op_func) {
    Dtype dtype = ret.dtype();
    size_t type_size = get_type_size(dtype);
    char* ret_data = static_cast<char*>(ret.data_ptr());
    size_t size = ret.size();

    RunParallel(static_cast<int>(size), [&](int i) {
        op_func(static_cast<void*>(ret_data + i * type_size), static_cast<const void*>(ret_data + i * type_size));
    });
}

// 一元操作
void ApplyOpSimple(NdArray& ret, const NdArray& src, UnaryElementOp op_func) {
    if (ret.shape() != src.shape()) {
        throw std::runtime_error("Shapes must match for unary operation");
    }

    Dtype dtype = ret.dtype();
    if (dtype != src.dtype()) {
        throw std::runtime_error("Data types must match for unary operation");
    }

    size_t type_size = get_type_size(dtype);
    char* ret_data = static_cast<char*>(ret.data_ptr());
    const char* src_data = static_cast<const char*>(src.data_ptr());
    size_t size = ret.size();

    RunParallel(static_cast<int>(size), [&](int i) {
        op_func(static_cast<void*>(ret_data + i * type_size), static_cast<const void*>(src_data + i * type_size));
    });
}

// 二元操作
void ApplyOpSimple(NdArray& ret, const NdArray& lhs, const NdArray& rhs, BinaryElementOp op_func, bool check_dtype) {
    if (ret.shape() != lhs.shape() || ret.shape() != rhs.shape()) {
        throw std::runtime_error("Shapes must match for binary operation");
    }

    Dtype dtype = ret.dtype();
    if (check_dtype) {
        if (dtype != lhs.dtype() || dtype != rhs.dtype()) {
            throw std::runtime_error("Data types must match for binary operation");
        }
    }

    size_t type_size = get_type_size(dtype);
    char* ret_data = static_cast<char*>(ret.data_ptr());
    const char* l_data = static_cast<const char*>(lhs.data_ptr());
    const char* r_data = static_cast<const char*>(rhs.data_ptr());
    size_t size = ret.size();

    RunParallel(static_cast<int>(size), [&](int i) {
        op_func(static_cast<void*>(ret_data + i * type_size), static_cast<const void*>(l_data + i * type_size),
                static_cast<const void*>(r_data + i * type_size));
    });
}

}  // namespace tinyndarray
